#include "nvds_pipeline.h"
#include "nvds_custom_meta.h"
#include "nvds_pipeline_builder.h"
#include "nvds_sourcebin.h"
#include <memory>
#include <spdlog/spdlog.h>
#include <gst-nvmessage.h>
#include <nvbufsurface.h>
#include <nvbufsurftransform.h>
#include <cuda_runtime_api.h>
#include <stdexcept>
#include <yaml-cpp/yaml.h>

static uint64_t getCurrentMilliseconds()
{
    auto tp = std::chrono::high_resolution_clock::now().time_since_epoch();
    return std::chrono::duration_cast<std::chrono::milliseconds>(tp).count();
}

NVDSPipeline::NVDSPipeline() : m_useDisplay(false), m_useJpgenc(false), m_latencyInfo(nullptr)
{
    gst_init(nullptr, nullptr);

    g_mutex_init(&m_latencyLock);
}

NVDSPipeline::~NVDSPipeline()
{
    g_mutex_clear(&m_latencyLock);

    gst_deinit();
}

void NVDSPipeline::start(const std::string &conf_file)
{
    if (!loadConfig(conf_file)) {
        throw std::runtime_error("配置文件读取失败");
    }

    int num_sources = m_videoUrls.size();
    if (m_latencyInfo == NULL) {
        m_latencyInfo = (NvDsFrameLatencyInfo *)calloc(1, num_sources * sizeof(NvDsFrameLatencyInfo));
    }

    // pipeline 初始化
    pipelineInit(m_videoUrls);
    GstBus *bus  = gst_pipeline_get_bus(GST_PIPELINE(m_pipeline));
    m_busWatchId = gst_bus_add_watch(bus, bus_call, m_gstLoop);
    gst_object_unref(bus);

    m_gstLoop = g_main_loop_new(NULL, FALSE);

    /* Set the pipeline to "playing" state */
    gst_element_set_state(m_pipeline, GST_STATE_PLAYING);

    /* Wait till pipeline encounters an error or EOS */
    g_main_loop_run(m_gstLoop);
}

void NVDSPipeline::stop()
{
    /* Set the pipeline to "null" state */
    gst_element_set_state(m_pipeline, GST_STATE_NULL);

    g_main_loop_unref(m_gstLoop);

    g_source_remove(m_busWatchId);

    // 释放 latency
    if (m_latencyInfo != NULL) {
        free(m_latencyInfo);
        m_latencyInfo = NULL;
    }

    // pipeline 销毁
    pipelineDestroy();
}

bool NVDSPipeline::loadConfig(const std::string &filename)
{
    if (m_videoUrls.size() > 0)
        m_videoUrls.clear();

    try {
        YAML::Node config = YAML::LoadFile(filename);
        for (auto video_url : config["source-list"]["list"]) {
            m_videoUrls.emplace_back(video_url.as<std::string>());
        }
        if (config["pipeline"].IsDefined()) {
            if (config["pipeline"]["display"].IsDefined())
                m_useDisplay = config["pipeline"]["display"].as<bool>();

            if (config["pipeline"]["jpeg_enc"].IsDefined()) {
                const YAML::Node &jpeg_enc = config["pipeline"]["jpeg_enc"];
                if (jpeg_enc["enable"].IsDefined())
                    m_useJpgenc = jpeg_enc["enable"].as<bool>();
                if (jpeg_enc["img_quality"].IsDefined())
                    m_jpgencParam.jpg_quality = jpeg_enc["img_quality"].as<int>();
                if (jpeg_enc["scale"].IsDefined()) {
                    const YAML::Node &scale_node = jpeg_enc["scale"];
                    if (scale_node["enable"].IsDefined())
                        m_jpgencParam.scale_enable = scale_node["enable"].as<bool>();
                    if (scale_node["width"].IsDefined())
                        m_jpgencParam.scale_width = scale_node["width"].as<int>();
                    if (scale_node["height"].IsDefined())
                        m_jpgencParam.scale_height = scale_node["height"].as<int>();
                }
            }
        }
    } catch (const std::exception &e) {
        spdlog::error("Error load config file {}.", filename);
        return false;
    }
    return true;
}

void NVDSPipeline::pipelineInit(const std::vector<std::string> &urls)
{
    m_streammux = gst_element_factory_make("nvstreammux", "stream-muxer");
    if (!m_streammux) {
        throw std::runtime_error("Gstreamer streammux create failed.");
    }

    // 创建 pipeline
    if (m_useDisplay) {
        m_pipeline = GstPipelineBuilder::buildPipelineWithDisplay("ds-pipeline",
                                                                  m_streammux,
                                                                  "deepstream_config.yml",
                                                                  urls.size(),
                                                                  pgie_src_pad_buffer_probe,
                                                                  this,
                                                                  tailer_sink_pad_buffer_probe,
                                                                  this);
    } else {
        m_pipeline = GstPipelineBuilder::buildPipeline("ds-pipeline",
                                                       m_streammux,
                                                       "deepstream_config.yml",
                                                       urls.size(),
                                                       pgie_src_pad_buffer_probe,
                                                       this,
                                                       tailer_sink_pad_buffer_probe,
                                                       this);
    }
    if (!m_pipeline) {
        throw std::runtime_error("Gstreamer pipeline create failed.");
    }

    // 获取进入检测模块的输入图像尺寸
    g_object_get(G_OBJECT(m_streammux), "width", &m_inputWidth, "height", &m_inputHeight, NULL);

    // 设置 sourcebin
    for (int i = 0; i < urls.size(); i++) {
        gchar bin_name[16] = {}, pad_name[16] = {};
        g_snprintf(bin_name, 15, "source-bin-%02d", i);
        g_snprintf(pad_name, 15, "sink_%u", i);

        auto sourcebin = std::make_shared<NVDSSourcebin>(bin_name,
                                                         urls[i].c_str(),
                                                         false,
                                                         m_jpgencParam.jpg_quality,
                                                         m_jpgencParam.scale_enable,
                                                         m_jpgencParam.scale_width,
                                                         m_jpgencParam.scale_height);
        sourcebin->bindPipeline(m_pipeline);
        sourcebin->bindStreammux(pad_name, m_streammux);
        m_sourcebins.push_back(sourcebin);
    }

    int current_device = -1;
    cudaGetDevice(&current_device);
    struct cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, current_device);
    m_objCtxHandle = nvds_obj_enc_create_context(current_device);
    if (!m_objCtxHandle) {
        throw std::runtime_error("Unable to create context!");
    }
}

void NVDSPipeline::pipelineDestroy()
{
    gst_object_unref(GST_OBJECT(m_pipeline));
    if (m_objCtxHandle)
        nvds_obj_enc_destroy_context(m_objCtxHandle);
    m_sourcebins.clear();
}

gboolean NVDSPipeline::bus_call(GstBus *bus, GstMessage *msg, gpointer data)
{
    GMainLoop *loop = (GMainLoop *)data;
    if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_EOS) {
        spdlog::error("End of stream.");
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_ERROR) {
        gchar  *debug;
        GError *error;
        gst_message_parse_error(msg, &error, &debug);
        spdlog::error("ERROR from element[{}] {}: {}.", error->code, GST_OBJECT_NAME(msg->src), error->message);
        if (debug)
            spdlog::error("Error details: {}.", debug);
        g_free(debug);
        g_clear_error(&error);
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_WARNING) {
        gchar  *debug;
        GError *error;
        gst_message_parse_warning(msg, &error, &debug);
        spdlog::warn("WARNING from element[{}] {}: {}.", error->code, GST_OBJECT_NAME(msg->src), error->message);
        if (debug)
            spdlog::warn("Warning details: {}.", debug);
        g_free(debug);
        g_error_free(error);
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_TAG) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_STATE_CHANGED) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_NEW_CLOCK) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_STREAM_STATUS) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_ELEMENT) {
        if (gst_nvmessage_is_stream_eos(msg)) {
            guint stream_id;
            if (gst_nvmessage_parse_stream_eos(msg, &stream_id)) {
                spdlog::warn("Got EOS from stream {}.", stream_id);
            }
        }
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_LATENCY) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_ASYNC_DONE) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_QOS) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_PROGRESS) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_STREAM_START) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_NEED_CONTEXT) {
    } else if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_HAVE_CONTEXT) {
    } else {
        spdlog::warn("Got unknown error code: {}.", int(GST_MESSAGE_TYPE(msg)));
    }
    return TRUE;
}

GstPadProbeReturn NVDSPipeline::pgie_src_pad_buffer_probe(GstPad *pad, GstPadProbeInfo *info, gpointer u_data)
{
    uint64_t      curtime = getCurrentMilliseconds();
    NVDSPipeline *pthis   = (NVDSPipeline *)u_data;
    GstBuffer    *buf     = GST_PAD_PROBE_INFO_BUFFER(info);
    NvBufSurface *surface = NULL;

    // nvidia jpg encode 功能
    if (pthis->m_useJpgenc) {
        GstMapInfo map_info = GST_MAP_INFO_INIT;
        if (gst_buffer_map(buf, &map_info, GST_MAP_READ)) {
            surface = (NvBufSurface *)map_info.data;
        }
        gst_buffer_unmap(buf, &map_info);
    }

    // user metadata 数据赋值
    NvDsBatchMeta *batch_meta = gst_buffer_get_nvds_batch_meta(buf);
    for (NvDsMetaList *frame = batch_meta->frame_meta_list; frame != NULL; frame = frame->next) {
        NvDsFrameMeta *frame_meta = (NvDsFrameMeta *)(frame->data);

        // 时间戳赋值
        for (NvDsUserMetaList *user_meta_list = frame_meta->frame_user_meta_list; user_meta_list != NULL;
             user_meta_list                   = user_meta_list->next) {
            NvDsUserMeta *user_meta = (NvDsUserMeta *)user_meta_list->data;
            if (user_meta->base_meta.meta_type == NVDSCustomMeta::getType()) {
                NVDSCustomMeta *user_metadata              = (NVDSCustomMeta *)user_meta->user_meta_data;
                user_metadata->timestamp()[DETECTION_TIME] = curtime;
            }
        }

        // 原始图片jpeg编码（无法使用图片缩放功能）
        if (pthis->m_useJpgenc && pthis->m_objCtxHandle) {
            NvDsObjEncUsrArgs frame_data = {0};
            frame_data.isFrame           = 1;
            frame_data.saveImg           = FALSE;
            frame_data.attachUsrMeta     = TRUE;
            frame_data.scaleImg          = FALSE;
            frame_data.scaledWidth       = 0;
            frame_data.scaledHeight      = 0;
            frame_data.quality           = pthis->m_jpgencParam.jpg_quality;
            nvds_obj_enc_process(pthis->m_objCtxHandle, &frame_data, surface, NULL, frame_meta);
        }
    }

    if (pthis->m_useJpgenc && pthis->m_objCtxHandle) {
        // 等待编码完成
        nvds_obj_enc_finish(pthis->m_objCtxHandle);
    }

    return GST_PAD_PROBE_OK;
}

GstPadProbeReturn NVDSPipeline::tailer_sink_pad_buffer_probe(GstPad *pad, GstPadProbeInfo *info, gpointer u_data)
{
    uint64_t      curtime = getCurrentMilliseconds();
    NVDSPipeline *pthis   = (NVDSPipeline *)u_data;
    GstBuffer    *buf     = GST_PAD_PROBE_INFO_BUFFER(info);

    // 打印时延统计信息
    if (nvds_enable_latency_measurement) {
        g_mutex_lock(&pthis->m_latencyLock);
        NvDsFrameLatencyInfo *latency_info         = pthis->m_latencyInfo;
        guint                 num_sources_in_batch = nvds_measure_buffer_latency(buf, latency_info);
        spdlog::info("************ source batch in num: {} ************", num_sources_in_batch);
        for (guint i = 0; i < num_sources_in_batch; i++) {
            spdlog::info("source id = {}, frame_num = {}, latency = {}(ms) ",
                         latency_info[i].source_id,
                         latency_info[i].frame_num,
                         latency_info[i].latency);
        }
        g_mutex_unlock(&pthis->m_latencyLock);
    }

    DetectResultList resultList;
    NvDsBatchMeta   *batch_meta = gst_buffer_get_nvds_batch_meta(buf);
    for (NvDsMetaList *frame = batch_meta->frame_meta_list; frame != NULL; frame = frame->next) {
        NvDsFrameMeta *frame_meta = (NvDsFrameMeta *)(frame->data);

        // 结果数据赋值
        DetectResult result;
        result.frame_id       = frame_meta->frame_num;
        result.device_id      = frame_meta->source_id;
        result.ntp_timestamp  = frame_meta->ntp_timestamp;
        result.nvds_over_time = curtime;

        // 获取 meta 数据信息
        for (NvDsUserMetaList *user_meta_list = frame_meta->frame_user_meta_list; user_meta_list != NULL;
             user_meta_list                   = user_meta_list->next) {
            NvDsUserMeta *user_meta = (NvDsUserMeta *)user_meta_list->data;

            // 获取  nvds_obj_enc_process 输出的图像数据
            if (user_meta->base_meta.meta_type == NVDS_CROP_IMAGE_META) {
                if (result.image_data.data.size() > 0)
                    continue;

                NvDsObjEncOutParams *enc_jpeg_image = (NvDsObjEncOutParams *)user_meta->user_meta_data;
                result.image_data.data.assign(enc_jpeg_image->outBuffer, enc_jpeg_image->outBuffer + enc_jpeg_image->outLen);
                result.image_data.width  = pthis->m_inputWidth;
                result.image_data.height = pthis->m_inputHeight;
            }

            // 获取自定义 metadata 数据
            if (user_meta->base_meta.meta_type == NVDSCustomMeta::getType()) {
                NVDSCustomMeta *user_metadata = (NVDSCustomMeta *)user_meta->user_meta_data;
                result.landing_time           = user_metadata->timestamp()[LANDING_TIME];
                result.detection_time         = user_metadata->timestamp()[DETECTION_TIME];

                // 启用 surfaceTransform 获取图片
                if (!user_metadata->imgBuffer().empty()) {
                    result.image_data.data   = std::move(user_metadata->imgBuffer());
                    result.image_data.width  = user_metadata->getImageSize().width;
                    result.image_data.height = user_metadata->getImageSize().height;
                }
            }
        }

        // 获取目标信息
        for (NvDsMetaList *obj = frame_meta->obj_meta_list; obj != NULL; obj = obj->next) {
            NvDsObjectMeta *obj_meta = (NvDsObjectMeta *)(obj->data);

            Object object;
            object.id         = obj_meta->object_id;
            object.left       = int(obj_meta->detector_bbox_info.org_bbox_coords.left);
            object.top        = int(obj_meta->detector_bbox_info.org_bbox_coords.top);
            object.width      = int(obj_meta->detector_bbox_info.org_bbox_coords.width);
            object.height     = int(obj_meta->detector_bbox_info.org_bbox_coords.height);
            object.class_id   = obj_meta->class_id;
            object.class_name = obj_meta->obj_label;
            object.confidence = obj_meta->confidence;
            result.objects.emplace_back(std::move(object));
        }
        resultList.emplace_back(std::move(result));
    }

    uint64_t overtime = getCurrentMilliseconds();
    for (auto &ret : resultList) {
        ret.nvds_over_time = overtime;
    }

    if (pthis->m_retCallback)
        pthis->m_retCallback(resultList);

    return GST_PAD_PROBE_OK;
}
